#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
#

"""
闭包代码生成器（Lambda 和嵌套函数统一）

通过策略模式统一 85% 的重复代码，仅保留 15% 的差异（函数体生成）
"""

from typing import Any, Callable, Optional
from mcpy.ir import IRNode
from mcpy.type_system import get_cpp_type, TypeKind
from ..core.context import CodeGenContext


class ClosureGenerator:
    """闭包 Struct 生成器（Lambda + 嵌套函数统一）

    设计原则：
    1. 提取 85% 公共逻辑
    2. 通过策略模式处理 15% 差异（函数体生成）
    3. 保持向后兼容

    使用示例：
        # Lambda
        generator = ClosureGenerator(
            node=lambda_node,
            context=context,
            body_fn=lambda_body_generator
        )

        # 嵌套函数
        generator = ClosureGenerator(
            node=func_node,
            context=context,
            body_fn=nested_body_generator
        )

        struct_code = generator.generate()
    """

    def __init__(
        self,
        node: IRNode,
        context: CodeGenContext,
        body_fn: Callable[[Any, CodeGenContext], str],
        name_fn: Optional[Callable[[Any], str]] = None,
    ):
        """初始化生成器

        Args:
            node: IRLambda 或 IRFunctionNode
            context: 代码生成上下文
            body_fn: 函数体生成策略（唯一差异点）
            name_fn: functor 名称生成函数（可选，用于自定义命名）
        """
        self.node = node
        self.context = context
        self.body_fn = body_fn
        self.name_fn = name_fn

        # 统一属性访问
        self.params = node.params  # type: ignore
        self.captured_vars = node.captured_vars  # type: ignore

        self.captured_cpp_types = node.get_extension_param("captured_cpp_types") or {}
        self.capture_strategies = node.get_extension_param("capture_strategies") or {}

    # ==================== 公共 API ====================

    def generate(self) -> str:
        """生成完整的 struct 代码"""
        lines = []

        # 1. 生成函数签名
        sig_type = self._gen_signature()

        # 2. struct 头部（分成两行，与原格式一致）
        functor_name = self._get_name()
        lines.append(f"struct {functor_name} ")
        lines.append(f"     : public mcpy::lambda_base<{functor_name}, {sig_type}> {{")

        # 3. get_param_names() 静态方法
        if self.params:
            lines.append(self._gen_param_names())
            lines.append("")

        # 4. 捕获变量成员
        if self.captured_vars:
            lines.append(self._gen_captured_members())
            lines.append("")

        # 5. 构造函数（包含默认值处理）
        constructor_code = self._gen_constructor()
        if constructor_code:  # 只在需要时添加
            lines.append(constructor_code)
            lines.append("")

        # 6. operator() 方法
        lines.append(self._gen_operator())

        # 7. struct 结束
        lines.append("};")

        return "\n".join(lines)

    # ==================== 私有方法：签名生成 ====================

    def _gen_signature(self) -> str:
        """生成函数签名：Ret(Arg1, Arg2, ...)"""
        # 参数类型
        cpp_param_types = []
        for param in self.params:
            type_info = param.get_extension_param("inferred_type")
            if type_info and type_info.kind not in (TypeKind.ANY, TypeKind.UNKNOWN):
                cpp_type = get_cpp_type(type_info, self.node)
                cpp_param_types.append(cpp_type if cpp_type else "mc::variant")
            else:
                cpp_param_types.append("mc::variant")

        # 返回类型
        cpp_return_type = self.node.get_extension_param("cpp_return_type")
        if not cpp_return_type:
            func_type = self.node.get_extension_param("inferred_type")
            if func_type and func_type.return_type:
                if func_type.return_type.kind == TypeKind.FUNCTION:
                    cpp_return_type = "mc::variant"
                else:
                    cpp_return_type = get_cpp_type(func_type.return_type, self.node)
                    if not cpp_return_type:
                        cpp_return_type = "mc::variant"
            else:
                cpp_return_type = "mc::variant"

        # 构造签名
        if cpp_param_types:
            return f"{cpp_return_type}({', '.join(cpp_param_types)})"
        else:
            return f"{cpp_return_type}()"

    def _gen_param_names(self) -> str:
        """生成 get_param_names() 静态方法"""
        param_names = ", ".join([f'"{p.name}"' for p in self.params])
        return f"""    static constexpr std::array<const char*, {len(self.params)}> get_param_names() {{
        return {{{param_names}}};
    }}"""

    # ==================== 私有方法：成员变量 ====================

    def _gen_captured_members(self) -> str:
        """生成捕获变量成员

        注意：Lambda 和嵌套函数的捕获策略不同
        - Lambda: const_ref 存储为 const T&（引用成员）
        - 嵌套函数: const_ref 存储为 T（值成员）
        """
        lines = []
        for var_name in self.captured_vars:
            cpp_type = self.captured_cpp_types.get(var_name, "mc::variant")
            strategy = self.capture_strategies.get(var_name, "const_ref")

            if strategy == "shared_ref":
                member_type = f"std::shared_ptr<{cpp_type}>"
            elif strategy == "const_ref":
                # Lambda 和嵌套函数都使用 const T&
                member_type = f"const {cpp_type}&"
            else:  # value
                member_type = cpp_type

            lines.append(f"    {member_type} {var_name};")

        return "\n".join(lines)

    # ==================== 私有方法：构造函数 ====================

    def _gen_constructor(self) -> str:
        """生成构造函数（包含默认值处理）"""
        has_defaults = any(p.default_value is not None for p in self.params)

        if not self.captured_vars and not has_defaults:
            # 无捕获、无默认值：不需要构造函数（编译器自动生成）
            return ""

        # 生成构造函数
        lines = []
        functor_name = self._get_name()

        # 参数列表
        params = []
        init_list = []

        for var_name in self.captured_vars:
            cpp_type = self.captured_cpp_types.get(var_name, "mc::variant")
            strategy = self.capture_strategies.get(var_name, "const_ref")

            if strategy == "shared_ref":
                params.append(f"std::shared_ptr<{cpp_type}> {var_name}_")
                init_list.append(f"{var_name}({var_name}_)")
            elif strategy == "const_ref":
                params.append(f"const {cpp_type}& {var_name}_")
                init_list.append(f"{var_name}({var_name}_)")
            else:  # value
                params.append(f"{cpp_type} {var_name}_")
                init_list.append(f"{var_name}(std::move({var_name}_))")

        # 构造函数声明
        param_list = ", ".join(params) if params else ""
        lines.append(f"    {functor_name}({param_list})")

        if init_list:
            lines.append(f"        : {', '.join(init_list)} {{")
        else:
            lines.append("        {")

        # 默认值处理
        if has_defaults:
            lines.append(self._gen_defaults_init())

        lines.append("    }")

        return "\n".join(lines)

    def _gen_defaults_init(self) -> str:
        """生成默认值初始化代码"""
        from mcpy.codegen.core.generator import get_global_generator

        gen = get_global_generator()

        default_values = []
        first_default_idx = len(self.params)

        for i, param in enumerate(self.params):
            if param.default_value is not None:
                if i < first_default_idx:
                    first_default_idx = i

                # 生成默认值表达式
                param_type_info = param.get_extension_param("inferred_type")
                param_cpp_type = param.get_extension_param("cpp_type")

                default_context = self.context.create_child_context()
                if param_type_info:
                    default_context.set_extension_param(
                        "expected_type", param_type_info
                    )
                if param_cpp_type:
                    param.default_value.set_extension_param(
                        "cpp_type", param_cpp_type, replace_existing=True
                    )

                default_code = gen.generate_expression(
                    param.default_value, default_context
                )
                default_values.append(f"mc::variant({default_code})")

        defaults_init = ", ".join(default_values)
        return f"        set_defaults({{{defaults_init}}}, {first_default_idx});"

    # ==================== 私有方法：operator() ====================

    def _gen_operator(self) -> str:
        """生成 operator() 方法（委托函数体生成给策略）"""
        # 分析参数类型
        param_types = []
        for param in self.params:
            type_info = param.get_extension_param("inferred_type")
            param_types.append(type_info)

        all_types_known = all(
            t and t.kind not in (TypeKind.ANY, TypeKind.UNKNOWN) for t in param_types
        )

        # 无参数情况
        if len(self.params) == 0:
            body_code = self.body_fn(self.node, self.context)
            
            # 检查是否需要显式返回类型（async函数）
            cpp_return_type = self.node.get_extension_param("cpp_return_type")
            if cpp_return_type:
                # async函数需要显式返回类型
                return f"""    {cpp_return_type} operator()() const {{
{self._indent(body_code, 2)}
    }}"""
            else:
                return f"""    auto operator()() const {{
{self._indent(body_code, 2)}
    }}"""

        # 生成参数列表
        template_params = []
        func_params = []

        if all_types_known:
            # 强类型
            for param in self.params:
                cpp_type = param.get_cpp_type()
                func_params.append(f"{cpp_type} {param.name}")
        else:
            # 模板
            for i, param in enumerate(self.params):
                type_info = param_types[i]
                if type_info and type_info.kind not in (TypeKind.ANY, TypeKind.UNKNOWN):
                    cpp_type = get_cpp_type(type_info, self.node)
                    if cpp_type:
                        func_params.append(f"{cpp_type} {param.name}")
                    else:
                        template_params.append(f"typename T{i}")
                        func_params.append(f"T{i}&& {param.name}")
                else:
                    template_params.append(f"typename T{i}")
                    func_params.append(f"T{i}&& {param.name}")

        param_list = ", ".join(func_params)

        # 生成函数体（委托给策略）
        body_code = self.body_fn(self.node, self.context)

        # 组装 operator()
        if template_params:
            template_decl = ", ".join(template_params)
            return f"""    template<{template_decl}>
    auto operator()({param_list}) const {{
{self._indent(body_code, 2)}
    }}"""
        else:
            return f"""    auto operator()({param_list}) const {{
{self._indent(body_code, 2)}
    }}"""

    # ==================== 工具方法 ====================

    def _get_name(self) -> str:
        """获取 functor 名称"""
        if self.name_fn:
            return self.name_fn(self.node)

        return (
            self.node.functor_name  # type: ignore
            or self.node.get_extension_param("functor_name")
            or f"__closure_{id(self.node) & 0xFFFF}"
        )

    @staticmethod  # type: ignore
    def _indent(code: str, level: int) -> str:
        """代码缩进"""
        indent = "    " * level
        lines = code.split("\n")
        return "\n".join(indent + line if line.strip() else "" for line in lines)
